import math
import torch
from torch import nn

from gensim.models import word2vec
from gensim.models.keyedvectors import KeyedVectors
import jieba

import wordcloud
import numpy as np

class PositionalEncoding(nn.Module):
    "Implement the PE function."
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

class SentimentTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        encoder_layer = nn.TransformerEncoderLayer(d_model=300, nhead=6)
        self.pos = PositionalEncoding(300, 0.1)
        self.trans = nn.TransformerEncoder(encoder_layer, num_layers=6)
        self.l1 = nn.Linear(300, 4)
        self.cls = nn.Softmax()

    def forward(self, x):
        x = self.pos(x)
        x = self.trans(x)
        x = self.l1(x)
        x = x.sum(1) / x.size(1)
        x = self.cls(x)

        return x

def runTrans():
    model = SentimentTransformer()
    a = model(torch.rand(1, 22, 300)).detach().numpy()[0]
    wc = wordcloud.WordCloud(font_path = 'C:/Windows/Fonts/simfang.ttf', background_color='white')

    labels = ["喜悦", "愤怒", "厌恶", "低落"]
    f1 = {labels[i]: a[i] for i in range(4)}

    wc.generate_from_frequencies(f1)
    wc.to_file("./infer.png")

if __name__ == "__main__":
    runTrans()